from langgraph.graph import StateGraph, add_messages
from typing import TypedDict, Annotated, Dict, Any

from owl_ai.domain.graph.nodes.chat_llm_node import ChatLLMNode
from owl_ai.domain.graph.nodes.doc_extract_node import DocumentExtractNode
from owl_ai.domain.graph.nodes.knowledge_retrieve_node import RAGRetrieveNode
from owl_ai.domain.model.graph_entity import GraphConfigEntity


def update_object(state, node_output):
    return state['app']


def update_node_params(state: Dict[str, Any], node_output: Any) -> Dict[str, Any]:
    """
    自定义 reducer 函数，用于更新 node_params 字段
    """
    # 如果 node_params 不存在，则初始化为空字典
    if state is None:
        state = {}

    if node_output:
        # 将节点的输出添加到 node_params 中
        state.update(node_output)
    return state


class ChatFlowState(TypedDict):
    app: Annotated[Any, update_object]
    node_params: Annotated[dict, update_node_params]


class ChatFlowCompile:
    """
    ChatFlow图编译器
    """

    @classmethod
    def compile(cls, graph_config: GraphConfigEntity):
        """
        根据配置进行图编译
        Args:
            graph_config: 图配置

        Returns: 编译可执行的图

        """
        state_graph = StateGraph(ChatFlowState)

        graph = graph_config.graph

        # 先进行节点配置
        nodes = graph.get("nodes")
        for node_config in nodes:
            node_id = node_config.get('id')
            node_type = node_config.get('type')
            # 开始、结束节点是特殊节点，不用单独加入，在edge处理即可
            if node_type == 'start' or node_type == 'end':
                continue
            if node_type == 'llm':
                node = ChatLLMNode(node_config)
                state_graph.add_node(node_id, node)
            elif node_type == 'doc_extract':
                node = DocumentExtractNode(node_config)
                state_graph.add_node(node_id, node)
            elif node_type == 'rag_retrieve':
                node = RAGRetrieveNode(node_config)
                state_graph.add_node(node_id, node)

        # 再进行边的配置
        node_edges = graph.get("edges")
        for node_edge in node_edges:
            source_node_id = node_edge.get('source')
            dest_node_id = node_edge.get("dest")

            state_graph.add_edge(source_node_id, dest_node_id)

        return state_graph.compile()
